from models.bert_sentiment import BertSentiment
from data_utils.basic_data import load_train_val_dataset
from data_utils.bert_sentiment_data import get_train_val_data_loader, get_test_loader_by_df
import torch
from config import conf
import pandas as pd
from os.path import join, exists
import os
from glob import glob

DEVICE_ID = 1  # conf.get('gpu', 'device_id')


def test_senti():
    device = torch.device("cuda:%s" % (DEVICE_ID) if torch.cuda.is_available() else "cpu")
    _, val_dataset = load_train_val_dataset(split_ratio=0.8)
    test_loader = get_test_loader_by_df(val_dataset, device, 2, maxlen=400)
    _, _, tokenizer = get_train_val_data_loader(device, 4, shuffle=True, maxlen=400)
    bert_attn_model = BertSentiment(learning_rate=5e-5, version_id=1).to(device)
    version_id = len(glob(bert_attn_model.model_root + '/*version*'))
    bert_attn_model = BertSentiment(5e-5, version_id=version_id).to(device)
    bert_attn_model.load_best_model()
    results = list(bert_attn_model.inference(test_loader))

    save_path = join(bert_attn_model.model_save_dir, 'data', 'raw_val_rs.csv')
    pd.DataFrame(results, columns=['id', 'negative']).to_csv(
        save_path, index=False)
    print('saved ', save_path)


if __name__ == '__main__':
    test_senti()
